package com.liam.translation.websocket;

import com.alibaba.fastjson.JSON;
import com.alibaba.fastjson.JSONObject;
import com.liam.translation.xunfei.RtasrClientApp;
import com.liam.translation.xunfei.NewTranslationService;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import javax.websocket.*;
import javax.websocket.server.ServerEndpoint;
import java.io.IOException;
import java.io.PipedInputStream;
import java.io.PipedOutputStream;
import java.nio.ByteBuffer;
import java.security.SignatureException;
import java.util.Map;
import java.util.concurrent.ConcurrentHashMap;

/**
 * 语音识别WebSocket服务器
 * 处理前端WebSocket连接，转发语音识别结果
 * 
 * @author liam
 */
@ServerEndpoint("/speech-ws")
public class SpeechWebSocketServer {

    private static final Logger logger = LoggerFactory.getLogger(SpeechWebSocketServer.class);

    /**
     * 存储所有连接的会话
     */
    private static final Map<String, Session> SESSIONS = new ConcurrentHashMap<>();

    /**
     * 存储每个会话对应的语音识别处理器
     */
    private static final Map<String, SpeechRecognitionHandler> HANDLERS = new ConcurrentHashMap<>();

    /**
     * 存储每个会话的翻译方向
     * 默认为中译英
     */
    private static final Map<String, TranslationDirection> TRANSLATION_DIRECTIONS = new ConcurrentHashMap<>();

    /**
     * WebSocket连接建立时调用
     */
    @OnOpen
    public void onOpen(Session session) {
        String sessionId = session.getId();
        SESSIONS.put(sessionId, session);
        // 默认设置为中译英
        TRANSLATION_DIRECTIONS.put(sessionId, TranslationDirection.CHINESE_TO_ENGLISH);
        logger.info("新的WebSocket连接已建立: {}", sessionId);

        // 发送连接成功消息
        sendMessage(session, createResponse("连接成功", "system", null));
    }

    /**
     * 接收到WebSocket文本消息时调用
     */
    @OnMessage
    public void onMessage(String message, Session session) {
        String sessionId = session.getId();
        logger.info("收到来自会话 {} 的文本消息: {}", sessionId, message);

        try {
            // 解析前端发送的JSON消息
            JSONObject jsonMessage = JSON.parseObject(message);
            String action = jsonMessage.getString("action");

            if ("start".equals(action)) {
                // 开始录音识别
                startRecognition(session);
            } else if ("stop".equals(action)) {
                // 停止录音识别
                stopRecognition(session);
            } else if ("switchDirection".equals(action)) {
                // 切换翻译方向
                String direction = jsonMessage.getString("direction");
                switchTranslationDirection(session, direction);
            } else {
                logger.warn("未知的操作类型: {}", action);
                sendMessage(session, createResponse("未知的操作类型", "error", null));
            }
        } catch (Exception e) {
            logger.error("处理消息时发生错误", e);
            sendMessage(session, createResponse("服务器处理请求失败", "error", null));
        }
    }

    /**
     * 切换翻译方向
     */
    private void switchTranslationDirection(Session session, String direction) {
        String sessionId = session.getId();
        TranslationDirection newDirection;

        if ("cn2en".equals(direction)) {
            newDirection = TranslationDirection.CHINESE_TO_ENGLISH;
            logger.info("会话 {} 切换翻译方向为: 中译英", sessionId);
            sendMessage(session, createResponse("已切换为中译英模式", "system", null));
        } else if ("en2cn".equals(direction)) {
            newDirection = TranslationDirection.ENGLISH_TO_CHINESE;
            logger.info("会话 {} 切换翻译方向为: 英译中", sessionId);
            sendMessage(session, createResponse("已切换为英译中模式", "system", null));
        } else {
            logger.warn("无效的翻译方向: {}", direction);
            sendMessage(session, createResponse("无效的翻译方向", "error", null));
            return;
        }

        TRANSLATION_DIRECTIONS.put(sessionId, newDirection);
    }

    /**
     * 接收到WebSocket二进制消息时调用
     */
    @OnMessage
    public void onMessage(ByteBuffer data, Session session) {
        String sessionId = session.getId();

        // 获取当前会话的处理器
        SpeechRecognitionHandler handler = HANDLERS.get(sessionId);
        if (handler != null) {
            try {
                // 获取音频数据
                byte[] audioData = new byte[data.remaining()];
                data.get(audioData);

                // 写入音频数据到管道
                handler.writeAudioData(audioData);

                // 记录日志（避免频繁打印）
                if (Math.random() < 0.01) { // 只记录约1%的数据包
                    logger.debug("收到来自会话 {} 的二进制音频数据: {} 字节", sessionId, audioData.length);
                }
            } catch (Exception e) {
                logger.error("处理音频数据时发生错误", e);
                sendMessage(session, createResponse("处理音频数据失败", "error", null));
            }
        } else {
            logger.warn("收到音频数据但没有对应的处理器，会话ID: {}", sessionId);
        }
    }

    /**
     * WebSocket连接关闭时调用
     */
    @OnClose
    public void onClose(Session session) {
        String sessionId = session.getId();
        logger.info("WebSocket连接已关闭: {}", sessionId);

        // 停止并清理资源
        stopRecognition(session);
        SESSIONS.remove(sessionId);
        TRANSLATION_DIRECTIONS.remove(sessionId);
    }

    /**
     * WebSocket连接发生错误时调用
     */
    @OnError
    public void onError(Session session, Throwable error) {
        String sessionId = session.getId();
        logger.error("WebSocket连接发生错误: {}", sessionId, error);

        // 停止并清理资源
        stopRecognition(session);
        SESSIONS.remove(sessionId);
        TRANSLATION_DIRECTIONS.remove(sessionId);
    }

    /**
     * 开始语音识别
     */
    private void startRecognition(Session session) {
        String sessionId = session.getId();

        // 如果已经有识别任务在运行，先停止
        stopRecognition(session);

        try {
            // 创建语音识别处理器
            SpeechRecognitionHandler handler = new SpeechRecognitionHandler(session);
            HANDLERS.put(sessionId, handler);

            // 启动识别线程
            Thread recognitionThread = new Thread(handler);
            recognitionThread.setDaemon(true);
            recognitionThread.setName("Recognition-" + sessionId);
            recognitionThread.start();

            logger.info("会话 {} 的语音识别已启动", sessionId);

            // 发送成功消息给客户端
            sendMessage(session, createResponse("语音识别已启动，请开始说话", "system", null));

        } catch (Exception e) {
            logger.error("启动语音识别失败", e);
            sendMessage(session, createResponse("启动语音识别失败: " + e.getMessage(), "error", null));
        }
    }

    /**
     * 停止语音识别
     */
    private void stopRecognition(Session session) {
        String sessionId = session.getId();
        SpeechRecognitionHandler handler = HANDLERS.remove(sessionId);

        if (handler != null) {
            try {
                handler.stop();
                logger.info("会话 {} 的语音识别已停止", sessionId);
            } catch (Exception e) {
                logger.error("停止会话 {} 的语音识别时发生错误", sessionId, e);
            }
        }
    }

    /**
     * 发送WebSocket消息
     */
    public static void sendMessage(Session session, String message) {
        try {
            if (session.isOpen()) {
                session.getBasicRemote().sendText(message);
            }
        } catch (IOException e) {
            logger.error("发送WebSocket消息失败", e);
        }
    }

    /**
     * 创建响应JSON字符串
     */
    private String createResponse(String text, String type, String error) {
        JSONObject response = new JSONObject();
        response.put("text", text);
        response.put("type", type);
        if (error != null) {
            response.put("error", error);
        }
        return response.toJSONString();
    }

    /**
     * 语音识别处理器
     * 负责与RtasrClientApp交互，处理语音识别
     */
    private class SpeechRecognitionHandler implements Runnable {
        private final Session session;
        private volatile boolean running = true;
        private PipedOutputStream audioOutputStream;
        private PipedInputStream audioInputStream;
        private final String sessionId;

        public SpeechRecognitionHandler(Session session) throws IOException {
            this.session = session;
            this.sessionId = session.getId();
            // 增加缓冲区大小，提高数据传输稳定性
            this.audioInputStream = new PipedInputStream(32768); // 增加到32KB
            this.audioOutputStream = new PipedOutputStream(audioInputStream);
        }

        @Override
        public void run() {
            try {
                // 创建自定义的结果处理器
                RtasrResultHandler resultHandler = new RtasrResultHandler(session);

                logger.info("会话 {} 准备调用讯飞API进行语音识别", sessionId);

                // 调用RtasrClientApp进行语音识别
                RtasrClientApp.processWebSocketStream(audioInputStream, resultHandler);

                // 等待停止信号
                while (running) {
                    Thread.sleep(100);
                }
            } catch (SignatureException e) {
                logger.error("会话 {} 的API签名验证失败", sessionId, e);
                sendMessage(session, createResponse("API签名验证失败，请检查APPID和API_KEY配置", "error", null));
            } catch (Exception e) {
                logger.error("会话 {} 的语音识别处理异常", sessionId, e);
                sendMessage(session, createResponse("语音识别处理异常: " + e.getMessage(), "error", null));
            } finally {
                cleanup();
            }
        }

        /**
         * 写入音频数据
         */
        public void writeAudioData(byte[] data) throws IOException {
            if (audioOutputStream != null && running) {
                try {
                    // 直接写入数据，不进行额外处理，提高速度
                    audioOutputStream.write(data);
                    // 只在必要时刷新缓冲区，减少IO操作
                    if (data.length >= 1024) {
                        audioOutputStream.flush();
                    }
                } catch (IOException e) {
                    if (running) { // 只有在仍在运行时才记录错误
                        logger.error("会话 {} 写入音频数据失败", sessionId, e);
                        throw e; // 重新抛出异常让调用者处理
                    }
                }
            }
        }

        /**
         * 停止处理
         */
        public void stop() {
            if (running) {
                running = false;
                logger.info("会话 {} 正在停止语音识别...", sessionId);

                // 发送结束标识
                try {
                    RtasrClientApp.sendEnd();
                } catch (Exception e) {
                    logger.error("会话 {} 发送结束标识失败", sessionId, e);
                }

                cleanup();
            }
        }

        /**
         * 清理资源
         */
        private void cleanup() {
            try {
                if (audioOutputStream != null) {
                    try {
                        audioOutputStream.close();
                    } catch (IOException e) {
                        logger.error("会话 {} 关闭音频输出流失败", sessionId, e);
                    } finally {
                        audioOutputStream = null;
                    }
                }

                if (audioInputStream != null) {
                    try {
                        audioInputStream.close();
                    } catch (IOException e) {
                        logger.error("会话 {} 关闭音频输入流失败", sessionId, e);
                    } finally {
                        audioInputStream = null;
                    }
                }

                logger.info("会话 {} 的资源清理完成", sessionId);
            } catch (Exception e) {
                logger.error("会话 {} 清理资源时发生异常", sessionId, e);
            }
        }
    }

    /**
     * 讯飞语音识别结果处理器
     */
    public static class RtasrResultHandler {
        private final Session session;
        private final NewTranslationService translationService;

        public RtasrResultHandler(Session session) {
            this.session = session;
            this.translationService = NewTranslationService.getInstance();
        }

        /**
         * 处理最终结果
         */
        public void handleFinalResult(String result) {
            // 发送原始识别结果
            sendMessage(session, createResponse(result, "final", null));

            // 获取当前会话的翻译方向
            String sessionId = session.getId();
            TranslationDirection direction = TRANSLATION_DIRECTIONS.getOrDefault(
                    sessionId, TranslationDirection.CHINESE_TO_ENGLISH);

            // 翻译结果并发送
            try {
                String translatedText;
                if (direction == TranslationDirection.CHINESE_TO_ENGLISH) {
                    // 中译英
                    translatedText = translationService.translateToEnglish(result);
                } else {
                    // 英译中
                    translatedText = translationService.translateToChineseFromEnglish(result);
                }

                sendMessage(session, createResponse(translatedText, "final-translation", null));
                logger.info("翻译成功 ({}): {} -> {}", direction, result, translatedText);
            } catch (Exception e) {
                logger.error("翻译失败", e);
                sendMessage(session, createResponse("翻译失败: " + e.getMessage(), "error", null));
            }
        }

        /**
         * 处理中间结果
         */
        public void handleInterimResult(String result) {
            // 发送原始中间结果
            sendMessage(session, createResponse(result, "interim", null));

            // 获取当前会话的翻译方向
            String sessionId = session.getId();
            TranslationDirection direction = TRANSLATION_DIRECTIONS.getOrDefault(
                    sessionId, TranslationDirection.CHINESE_TO_ENGLISH);

            // 翻译中间结果并发送
            try {
                // 只有当结果长度超过一定值才进行翻译，避免频繁调用API
                if (result.length() > 5) {
                    String translatedText;
                    if (direction == TranslationDirection.CHINESE_TO_ENGLISH) {
                        // 中译英
                        translatedText = translationService.translateToEnglish(result);
                    } else {
                        // 英译中
                        translatedText = translationService.translateToChineseFromEnglish(result);
                    }

                    sendMessage(session, createResponse(translatedText, "interim-translation", null));
                }
            } catch (Exception e) {
                // 中间结果翻译失败不发送错误消息，只记录日志
                logger.warn("中间结果翻译失败: {}", e.getMessage());
            }
        }

        /**
         * 处理错误
         */
        public void handleError(String errorMessage) {
            sendMessage(session, createResponse(null, "error", errorMessage));
        }

        /**
         * 创建响应JSON字符串
         */
        private String createResponse(String text, String type, String error) {
            JSONObject response = new JSONObject();
            response.put("text", text);
            response.put("type", type);
            if (error != null) {
                response.put("error", error);
            }
            return response.toJSONString();
        }
    }

    /**
     * 翻译方向枚举
     */
    public enum TranslationDirection {
        CHINESE_TO_ENGLISH("cn2en", "中译英"),
        ENGLISH_TO_CHINESE("en2cn", "英译中");

        private final String code;
        private final String description;

        TranslationDirection(String code, String description) {
            this.code = code;
            this.description = description;
        }

        public String getCode() {
            return code;
        }

        public String getDescription() {
            return description;
        }
    }
}